Implement SystemParam for SmallVec

[Shatur]

Objective

To register a system from a script, I need to support for dynamic number of system parameters (since scripting languages usually don't play well with generics). Right now the only solution is to use Vec<T> as a system param which is allocated every time the system runs. But it would be nice to avoid allocations when the number of system parameters is low.

Solution

Implement SystemParam for SmallVec. It's already a dependency for bevy_ecs and allows to avoid allocations with user-defined number of system parameters.

Testing

• Did you test these changes? If so, how? I tried creating a system with a SmallVec<[Local<usize>; 8]>.
• Are there any parts that need more testing? I don't think so. The code follows Vec impl.

[Shatur]

Looks like CI failure is unrelated.

[kfc35]

Looks good to me, follows the regular Vec impl except for one thing: the unsafe impl for ParamSet<'_, '_, SmallVec<[T; N]>>. Is that also necessary to implement? I assume since it works for you, it doesn’t need to be done but just figured I’d ask the question (for my learning)

[Shatur]

Is that also necessary to implement?

I personally don't need it, just implemented it for consistency (since Vec has it.)

[chescock]

Oh, yay, I'm happy to see SystemParamBuilder is getting use!

But it would be nice to avoid allocations when the number of system parameters is low.

Hmm, I wonder whether we could even make a Vec-like parameter that avoids per-run allocations entirely by re-using a buffer between runs.

[chescock]

I think you also need an impl<...> SystemParamBuilder<SmallVec<[P; N]>> for SmallVec<[B; N]> in builder.rs so that you can actually configure the systems.

[Shatur]

Ah, right, I didn't test the builder yet 😅 Added!

[chescock]

It looks like you copied the impl block for ParamSet<Vec<_>>, but not the impl SystemParam, so there's no way to actually obtain a ParamSet<SmallVec>.

I don't think you need a SmallVec version of ParamSet, though. ParamSet<Vec> never actually allocates a Vec, and just creates the subparameters on-demand in get_mut, so there is no allocation to save. So I'd be inclined to delete this impl block.

[Shatur]

Makes sense, removed!

[Shatur]

Hmm, I wonder whether we could even make a Vec-like parameter that avoids per-run allocations entirely by re-using a buffer between runs.

This was my original goal, but I couldn't find where to store the Vec, so I went with SmallVec. If you have any idea, I'm happy to try.

[chescock]

Hmm, I wonder whether we could even make a Vec-like parameter that avoids per-run allocations entirely by re-using a buffer between runs.

This was my original goal, but I couldn't find where to store the Vec, so I went with SmallVec. If you have any idea, I'm happy to try.

I did a few experiments, but without much success.

The "buffer reuse" technique from Wild Performance Tricks seems like it should work, but if you try to put a Vec<T::Item<'static, 'static>> in the state then the compiler complains that it can't prove Self::Item::State == Self::State because it can't prove T::Item<'w, 's>::Item<'static, 'static> == T::Item<'static, 'static>.

Then I thought that a Vec that dropped its contents but didn't deallocate sounded like bumpalo::collections::Vec, which we already have a dependency on! But the Item type would be Vec<'state, T::Item<'world, 'state>> and that's invalid because it requires T::Item<'world, 'state>: 'state, which doesn't hold.

I think I might have gotten it working using raw pointers! You can use alloc::alloc() to manually allocate a buffer and hold the raw pointer in the State. It's a lot of unsafe, though.

[Shatur]

I think I might have gotten it working using raw pointers! You can use alloc::alloc() to manually allocate a buffer and hold the raw pointer in the State. It's a lot of unsafe, though.

Would be interesting to see!

[chescock]

Would be interesting to see!

Alright, I got it building and passing at least a simple test in miri, so I think this might work? It needs a lot of polish, including safety comments and actual names, but I think it could work. I'm not sure where else to post it, so here it is:

Details

/// A slice that owns its values but not the underlying allocation.
pub struct SomethingVec<'a, T> {
    ptr: *mut [T],
    // This is logically `&'a mut [T]`, but we don't want to require `T: 'a`
    marker: PhantomData<(&'a mut (), [T])>,
}

impl<T> Drop for SomethingVec<'_, T> {
    fn drop(&mut self) {
        unsafe { drop_in_place(self.ptr) };
    }
}

// Raw pointers disable `Sync` and `Send`,
// but this should act as `[T]`.
unsafe impl<T: Sync> Sync for SomethingVec<'_, T> {}
unsafe impl<T: Send> Send for SomethingVec<'_, T> {}

impl<T> Deref for SomethingVec<'_, T> {
    type Target = [T];

    fn deref(&self) -> &Self::Target {
        unsafe { &*self.ptr }
    }
}

impl<T> DerefMut for SomethingVec<'_, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut *self.ptr }
    }
}

// TODO: More traits!  See what `Box<[T]>` implements.
// `IntoIterator` is especially important,
// and probably needs a custom `IntoIter` type.

#[doc(hidden)]
pub struct SomethingVecState<S> {
    states: Vec<S>,
    /// A pointer to an owned buffer big enough to hold
    /// a slice of `states.len()` `SystemParam::Item` values.
    ptr: *mut u8,
    /// The layout used to allocate `ptr`.
    ///
    /// We have to erase the actual `SystemParam` type to erase
    /// its lifetimes, so we cannot reconstruct this from types
    /// in the `Drop` impl.`
    layout: Layout,
}

impl<S> Drop for SomethingVecState<S> {
    fn drop(&mut self) {
        if self.layout.size() != 0 {
            unsafe { alloc::alloc::dealloc(self.ptr, self.layout) };
        }
    }
}

// Raw pointers disable `Sync` and `Send`,
// but the data behind `ptr` is ununitialized
// except when borrowed by `SomethingVec`.
unsafe impl<S: Sync> Sync for SomethingVecState<S> {}
unsafe impl<S: Send> Send for SomethingVecState<S> {}

/// A [`SystemParamBuilder`] for a [`SomethingVec`].
pub struct SomethingVecBuilder<B>(pub Vec<B>);

unsafe impl<P: SystemParam, B: SystemParamBuilder<P>> SystemParamBuilder<SomethingVec<'_, P>>
    for SomethingVecBuilder<B>
{
    fn build(self, world: &mut World) -> <SomethingVec<'_, P> as SystemParam>::State {
        // Note that we use `P::Item<'static, 'static>` to match `P::Item<'w, 's>` later.
        // These will be the same layout as `P` in practice, but it's technically possible
        // for a custom `SystemParam` to violate that.
        let layout = Layout::array::<P::Item<'static, 'static>>(self.0.len()).unwrap();
        let ptr = if layout.size() != 0 {
            unsafe { alloc(layout) }
        } else {
            dangling_mut::<P::Item<'static, 'static>>().cast()
        };
        let states = self
            .0
            .into_iter()
            .map(|builder| builder.build(world))
            .collect();
        SomethingVecState {
            states,
            layout,
            ptr,
        }
    }
}

unsafe impl<T: SystemParam> SystemParam for SomethingVec<'_, T> {
    type State = SomethingVecState<T::State>;

    type Item<'world, 'state> = SomethingVec<'state, T::Item<'world, 'state>>;

    fn init_state(_world: &mut World) -> Self::State {
        SomethingVecState {
            states: Vec::new(),
            layout: Layout::new::<[T::Item<'static, 'static>; 0]>(),
            ptr: dangling_mut::<T::Item<'static, 'static>>().cast(),
        }
    }

    fn init_access(
        state: &Self::State,
        system_meta: &mut SystemMeta,
        component_access_set: &mut FilteredAccessSet,
        world: &mut World,
    ) {
        for state in &state.states {
            T::init_access(state, system_meta, component_access_set, world);
        }
    }

    unsafe fn get_param<'world, 'state>(
        state: &'state mut Self::State,
        system_meta: &SystemMeta,
        world: UnsafeWorldCell<'world>,
        change_tick: Tick,
    ) -> Result<Self::Item<'world, 'state>, SystemParamValidationError> {
        // The layout was allocated as `T::Item<'static, 'static>`,
        // and layouts cannot depend on lifetimes, so the layout will be valid.
        debug_assert_eq!(
            state.layout,
            Layout::array::<T::Item<'world, 'state>>(state.states.len()).unwrap()
        );
        // We have `&'state mut Self::State`, so nothing else can access the
        // memory pointed to by `ptr` for the duration of `'state`.
        let ptr = state.ptr.cast::<T::Item<'_, '_>>();
        let mut len = 0;
        // Keep a valid `SomethingVec` updated with the current length
        // so that we drop any existing params if we return early.
        let mut result = SomethingVec {
            ptr: slice_from_raw_parts_mut(ptr, len),
            marker: PhantomData,
        };
        for state in &mut state.states {
            let param = unsafe { T::get_param(state, system_meta, world, change_tick) }?;
            unsafe { ptr.add(len).write(param) };
            len += 1;
            result.ptr = slice_from_raw_parts_mut(ptr, len);
        }
        Ok(result)
    }

    fn apply(state: &mut Self::State, system_meta: &SystemMeta, world: &mut World) {
        for state in &mut state.states {
            T::apply(state, system_meta, world);
        }
    }

    fn queue(state: &mut Self::State, system_meta: &SystemMeta, mut world: DeferredWorld) {
        for state in &mut state.states {
            T::queue(state, system_meta, world.reborrow());
        }
    }
}

[Shatur]

You didn't joke about the amount of unsafe 😅

[Shatur]

I added this to the 0.19 milestone since it's a very trivial non-breaking addition and it would be nice to have this in for my upcoming scripting crate 🙂